# -*- coding: utf-8 -*-
"""
Web应用数据库管理

为FastAPI提供异步数据库支持
"""

import logging
from typing import AsyncGenerator
from sqlalchemy.ext.asyncio import AsyncSession, create_async_engine, async_sessionmaker
from sqlalchemy.exc import SQLAlchemyError

from .config import get_settings

logger = logging.getLogger(__name__)
settings = get_settings()

# 全局变量
engine = None
async_session_factory = None


async def init_database():
    """初始化数据库"""
    global engine, async_session_factory
    
    try:
        # 创建异步引擎
        # SQLite不支持连接池参数，需要区别处理
        engine_kwargs = {"echo": settings.DEBUG}
        
        if not settings.DATABASE_URL.startswith("sqlite"):
            # 非SQLite数据库添加连接池参数
            engine_kwargs.update({
                "pool_size": settings.DATABASE_POOL_SIZE,
                "max_overflow": settings.DATABASE_MAX_OVERFLOW
            })
        
        engine = create_async_engine(settings.DATABASE_URL, **engine_kwargs)
        
        # 创建会话工厂
        async_session_factory = async_sessionmaker(
            engine,
            class_=AsyncSession,
            expire_on_commit=False
        )
        
        logger.info("数据库初始化完成")
        
    except Exception as e:
        logger.error(f"数据库初始化失败: {e}")
        raise


async def close_database():
    """关闭数据库连接"""
    global engine
    
    if engine:
        await engine.dispose()
        logger.info("数据库连接已关闭")


async def get_db() -> AsyncGenerator[AsyncSession, None]:
    """获取数据库会话"""
    if not async_session_factory:
        raise RuntimeError("数据库未初始化")
    
    async with async_session_factory() as session:
        try:
            yield session
        except SQLAlchemyError as e:
            await session.rollback()
            logger.error(f"数据库操作异常: {e}")
            raise
        finally:
            await session.close()